import numpy as np
from pyevtk.hl import gridToVTK
import time


class GDS:  # 定义装药数据结构类GDS(Grain Data Structure)
    def __init__(self, lx: float = 1.0,  # 立方体的长宽高
                 ly: float = 1.0,
                 lz: float = 1.0,
                 nx_points: int = 50,  # 边上的节点数
                 ny_points: int = 50,
                 nz_points: int = 50,
                 n_slots: int = 6,
                 m: float = 2,
                 rho_refine: float = 1.0,  # 柱坐标网格加密比例
                 theta_refine: float = 1.0,
                 zeta_refine: float = 1.0):
        # 存储为成员变量
        self.lx, self.ly, self.lz = lx, ly, lz
        self.nx_points, self.ny_points, self.nz_points = nx_points, ny_points, nz_points
        self.n_slots = n_slots
        self.m = m
        self.rho_refine, self.theta_refine, self.zeta_refine = rho_refine, theta_refine, zeta_refine
        # 这里存放各种物理场变量，可被外界访问
        # —————————————————————————以下可扩展——————————————————————————————
        self.burning_surface_coordinate = None
        self.cavity_region_coordinate = None
        self.solid_region_coordinate = None
        self.cavity_region_sdf=None
        self.solid_region_sdf=None
        # —————————————————————————以上可扩展——————————————————————————————
        if self.nx_points!=1 or self.ny_points!=1 or self.nz_points!=1:
           self.refresh_grid()
        else:
            self.xyz_field=None
            self.xyz_field_cyl=None

    def refresh_grid(self):
        # 更新网格
        ########################### 笛卡尔坐标 #############################
        #             (0,ny_points-1,0)
        #             *─────────────────* (nx_points-1,ny_points,0)
        #            /│                /│
        #           / │               / │
        #          /  │y             /  ly
        #         /   ▲             /   │
        #        /    │   x        /    │
        #       /     *────►─lx───/─────*
        #      /     /           /     /
        #     /     /           /     /
        #    *─────────────────*     /
        #    │    /            │    lz
        #    │   /             │   /
        #    │  /              │  /
        #    │ /               │ /
        #    │/                │/
        #    *─────────────────* (nx_points-1,0,nz_points-1)
        #   /
        #  /
        # / Axis
        # 节点坐标(使用32位浮点数节省存储空间)
        self.x_vector = np.linspace(
            0, self.lx, self.nx_points, dtype=np.float32)  # x方向节点分布，向量
        self.y_vector = np.linspace(
            0, self.ly, self.ny_points, dtype=np.float32)  # y方向节点分布，向量
        self.z_vector = np.linspace(
            0, self.lz, self.nz_points, dtype=np.float32)  # z方向节点分布，向量
        x_field, y_field, z_field = np.meshgrid(
            # x场，y场，z场张量，注意：indexing务必选择'ij'
            self.x_vector, self.y_vector, self.z_vector, indexing='ij')

        # 形成结构网格下的坐标场
        # axis=0沿x方向，axis=1沿y方向，axis=2沿z方向，axis=3记录了(x,y,z)坐标
        # 这样的数据格式有利于直接作为神经网络(神经场)的输入
        self.xyz_field = np.stack((x_field,
                                   y_field,
                                   z_field), axis=-1)
        ########################### 圆柱坐标 #############################
        # 形成加密扇形结构网格下的坐标场(只包含装药所处的扇形拉伸体)，采用柱坐标Cylindrical coordinates
        self.rho_vector = np.linspace(
            # rho方向节点分布，向量
            0.1*self.lx, self.lx, int(self.nx_points*self.rho_refine), dtype=np.float32)
        self.theta_vector = np.linspace(
            # theta方向节点分布，向量
            0, np.pi/self.n_slots, int(self.ny_points*self.theta_refine), dtype=np.float32)
        self.zeta_vector = np.linspace(
            # zeta方向节点分布，向量
            0, self.lz, int(self.nz_points*self.zeta_refine), dtype=np.float32)

        rho_field, theta_field, zeta_field = np.meshgrid(
            # rho场，theta场，zeta场张量，注意：indexing务必选择'ij'
            self.rho_vector, self.theta_vector, self.zeta_vector, indexing='ij')
        x_field_cyl = rho_field*np.cos(theta_field)
        y_field_cyl = rho_field*np.sin(theta_field)
        z_field_cyl = zeta_field
        self.xyz_field_cyl = np.stack((x_field_cyl,
                                       y_field_cyl,
                                       z_field_cyl), axis=-1)

    def get_xyz_field(self):
        # 返回笛卡尔结构网格下的坐标场
        return self.xyz_field

    def get_zero_x_boundary(self):
        # 返回笛卡尔结构网格下x=0平面的坐标场
        return self.xyz_field[0]

    def get_zero_y_boundary(self):
        # 返回笛卡尔结构网格下y=0平面的坐标场
        return self.xyz_field[:, 0]

    def get_xyz_field_cyl(self):
        # 返回柱坐结构网格下的坐标场
        return self.xyz_field_cyl

    def get_bottom_theta_boundary(self):
        # 返回柱坐结构网格下theta=0平面的坐标场
        return self.xyz_field_cyl[:, 0]

    def get_top_theta_boundary(self):
        # 返回柱坐结构网格网格下theta=pi/n_slots平面的坐标场
        return self.xyz_field_cyl[:, -1]
    
    def get_inner_ring_boundary(self):
        # 返回柱坐结构网格网格下最内层环的坐标场
        return self.xyz_field_cyl[0]
    
    def get_outer_ring_boundary(self):
        # 返回柱坐结构网格网格下最外层环的坐标场
        return self.xyz_field_cyl[-1]
    

    def save_as_npz(self, npz_path: str):  # 保存为npz文件(速度最快，占用空间最小的保存方式)
        # 输出数组文件(可扩展)
        np.savez_compressed(npz_path+".npz",
                            lx=self.lx,  # x方向长度
                            ly=self.ly,  # y方向长度
                            lz=self.lz,  # z方向长度
                            nx_points=self.nx_points,  # x方向节点数
                            ny_points=self.ny_points,  # y方向节点数
                            nz_points=self.nz_points,  # z方向节点数
                            n_slots=self.n_slots,  # 装药槽数
                            m=self.m,  # 封头椭球系数
                            rho_refine=self.rho_refine,  # rho方向网格加密比例
                            theta_refine=self.theta_refine,  # theta方向网格加密比例
                            zeta_refine=self.zeta_refine,  # zeta方向网格加密比例
                            # —————————————————————————以下可扩展——————————————————————————————
                            burning_surface_coordinate=self.burning_surface_coordinate,  # 燃面坐标
                            cavity_region_coordinate=self.cavity_region_coordinate, # 装药空腔坐标
                            solid_region_coordinate=self.solid_region_coordinate, # 固体装药坐标
                            cavity_region_sdf=self.cavity_region_sdf, # 装药空腔SDF
                            solid_region_sdf=self.solid_region_sdf # 固体装药SDF
                            # —————————————————————————以上可扩展——————————————————————————————
                            )

    def save_as_vtk_in_cartesian(self, vts_path: str, field_dict: dict = {}):
        # 将笛卡尔网格保存为VTK文件(可视化最方便，但速度较慢，占用空间较大的保存方式)输出VTK文件，可以用Paraview查看
        x_field, y_field, z_field = np.meshgrid(
            # x场，y场，z场张量，注意：indexing务必选择'ij'
            self.x_vector, self.y_vector, self.z_vector, indexing='ij')
        # 输出VTK文件
        gridToVTK(
            vts_path,
            x_field,
            y_field,
            z_field,
            pointData=field_dict
        )

        
    def save_as_vtk_in_cyl(self, vts_path: str, field_dict: dict = {}):
        # 将柱坐标网格保存为VTK文件(可视化最方便，但速度较慢，占用空间较大的保存方式)输出VTK文件，可以用Paraview查看
        rho_field, theta_field, zeta_field = np.meshgrid(
        # rho场，theta场，zeta场张量，注意：indexing务必选择'ij'
        self.rho_vector, self.theta_vector, self.zeta_vector, indexing='ij')
        x_field_cyl = rho_field*np.cos(theta_field)
        y_field_cyl = rho_field*np.sin(theta_field)
        z_field_cyl = zeta_field
        # 输出VTK文件
        gridToVTK(
            vts_path,
            x_field_cyl,
            y_field_cyl,
            z_field_cyl,
            pointData=field_dict
        )


def load_npz(npz_path: str, need_compress: bool = False) -> GDS:
    # 读取npz文件,如果need_compress==False那么压缩nx_points，ny_points，nz_points为1
    npz_file = np.load(npz_path+".npz")
    if(need_compress):
        nx_points=1
        ny_points=1
        nz_points=1
    else:
        nx_points=npz_file["nx_points"]
        ny_points=npz_file["ny_points"]
        nz_points=npz_file["nz_points"]
    gds = GDS(lx=npz_file["lx"],
              ly=npz_file["ly"],
              lz=npz_file["lz"],
              nx_points=nx_points,
              ny_points=ny_points,
              nz_points=nz_points,
              n_slots=npz_file["n_slots"],
              m=npz_file["m"],
              rho_refine=npz_file["rho_refine"],
              theta_refine=npz_file["theta_refine"],
              zeta_refine=npz_file["zeta_refine"])
    # —————————————————————————以下可扩展——————————————————————————————
    # 燃面坐标
    gds.burning_surface_coordinate = npz_file["burning_surface_coordinate"]
    gds.cavity_region_coordinate= npz_file["cavity_region_coordinate"]
    gds.solid_region_coordinate = npz_file["solid_region_coordinate"]
    gds.cavity_region_sdf = npz_file["cavity_region_sdf"]
    gds.solid_region_sdf = npz_file["solid_region_sdf"]
    # —————————————————————————以上可扩展——————————————————————————————
    return gds


def save_as_vtk_from_npz(vts_path: str, npz_path: str):  # 将npz文件转存为vtk
    # 读取npz文件
    gds = load_npz(npz_path)
    # 输出VTK文件
    gds.save_as_vtk_in_cartesian(vts_path)
    gds.save_as_vtk_in_cyl(f"{vts_path}_cyl")
